from ultralytics import YOLO
import os
import math
import statistics
import matplotlib.pyplot as plt  # 导入 Matplotlib 库
from scipy.stats import norm
import numpy as np

# Load a model
model = YOLO("yolov8n-pose.pt")  # pretrained YOLOv8n model

# Specify the folder containing the images
image_folder = "D:/python/v8_mine/dataset/"
output_file = "D:/python/v8_mine/angles.txt"
bili_values = []
# 用于存储 closest_angle 的列表
closest_angles = []

# Open the output file for writing
with open(output_file, "w") as f:
    # Loop through each image file in the folder
    for filename in os.listdir(image_folder):
        if filename.endswith(".jpg") or filename.endswith(".png"):  # Adjust extensions as needed
            # Construct the full file path
            file_path = os.path.join(image_folder, filename)

            # Run inference on the image
            results = model(file_path)  # return a list of Results objects
            keypoints = results[0].keypoints

            # Transfer xy tensor to CPU if necessary
            # xy_cpu = keypoints.xy.cpu() if keypoints.xy.device != 'cpu' else keypoints.xy
            xyn_cpu = keypoints.xyn.cpu() if keypoints.xyn.device != 'cpu' else keypoints.xyn

            # Extract coordinates for right and left arm keypoints
            # shoulder_x_r, shoulder_y_r = xy_cpu[0, 6, 0].item(), xy_cpu[0, 6, 1].item()
            #elbow_x_r, elbow_y_r = xy_cpu[0, 8, 0].item(), xy_cpu[0, 8, 1].item()
            #wrist_x_r, wrist_y_r = xy_cpu[0, 10, 0].item(), xy_cpu[0, 10, 1].item()
            #shoulder_x_l, shoulder_y_l = xy_cpu[0, 5, 0].item(), xy_cpu[0, 5, 1].item()
            #elbow_x_l, elbow_y_l = xy_cpu[0, 7, 0].item(), xy_cpu[0, 7, 1].item()
            #wrist_x_l, wrist_y_l = xy_cpu[0, 9, 0].item(), xy_cpu[0, 9, 1].item()
            #hip_y_l = xy_cpu[0, 11, 1]
            #hip_y_r = xy_cpu[0, 12, 1]
            #eye_y_l = xy_cpu[0, 1, 1]
            #eye_y_r = xy_cpu[0, 2, 1]
            shoulder_x_r, shoulder_y_r = xyn_cpu[0, 6, 0].item(), xyn_cpu[0, 6, 1].item()
            elbow_x_r, elbow_y_r = xyn_cpu[0, 8, 0].item(), xyn_cpu[0, 8, 1].item()
            wrist_x_r, wrist_y_r = xyn_cpu[0, 10, 0].item(), xyn_cpu[0, 10, 1].item()
            shoulder_x_l, shoulder_y_l = xyn_cpu[0, 5, 0].item(), xyn_cpu[0, 5, 1].item()
            elbow_x_l, elbow_y_l = xyn_cpu[0, 7, 0].item(), xyn_cpu[0, 7, 1].item()
            wrist_x_l, wrist_y_l = xyn_cpu[0, 9, 0].item(), xyn_cpu[0, 9, 1].item()
            hip_y_l = xyn_cpu[0, 11, 1]
            hip_y_r = xyn_cpu[0, 12, 1]
            eye_y_l = xyn_cpu[0, 1, 1]
            eye_y_r = xyn_cpu[0, 2, 1]

            # Calculate vectors for right and left arms
            vector_A_x_r = elbow_x_r - shoulder_x_r
            vector_A_y_r = elbow_y_r - shoulder_y_r
            vector_B_x_r = wrist_x_r - elbow_x_r
            vector_B_y_r = wrist_y_r - elbow_y_r

            vector_A_x_l = elbow_x_l - shoulder_x_l
            vector_A_y_l = elbow_y_l - shoulder_y_l
            vector_B_x_l = wrist_x_l - elbow_x_l
            vector_B_y_l = wrist_y_l - elbow_y_l

            # Calculate dot products
            dot_product_r = vector_A_x_r * vector_B_x_r + vector_A_y_r * vector_B_y_r
            dot_product_l = vector_A_x_l * vector_B_x_l + vector_A_y_l * vector_B_y_l

            # Calculate magnitudes
            magnitude_A_r = math.sqrt(vector_A_x_r ** 2 + vector_A_y_r ** 2)
            magnitude_B_r = math.sqrt(vector_B_x_r ** 2 + vector_B_y_r ** 2)
            magnitude_A_l = math.sqrt(vector_A_x_l ** 2 + vector_A_y_l ** 2)
            magnitude_B_l = math.sqrt(vector_B_x_l ** 2 + vector_B_y_l ** 2)

            if magnitude_A_r * magnitude_B_r == 0:
                continue  # Skip this image if right arm magnitudes result in zero denominator
            if magnitude_A_l * magnitude_B_l == 0:
                continue  # Skip this image if left arm magnitudes result in zero denominator

            # Calculate cosine values
            cosine_angle_r = dot_product_r / (magnitude_A_r * magnitude_B_r)
            cosine_angle_l = dot_product_l / (magnitude_A_l * magnitude_B_l)

            # Calculate angles in degrees
            angle_in_degrees_r = math.degrees(math.acos(cosine_angle_r))
            angle_in_degrees_l = math.degrees(math.acos(cosine_angle_l))

            # Determine which angle is closest to 90 degrees
            closest_angle = min(angle_in_degrees_r, angle_in_degrees_l, key=lambda x: abs(x - 90))
            closest_angles.append(closest_angle)  # 将 closest_angle 添加到列表中

            #height_l = shoulder_y_l - hip_y_l
            #height_r = shoulder_y_r - hip_y_r

            #delta_y_hiptoshoulder = max(height_l, height_r)
            #delta_x_shoulder = math.sqrt((shoulder_x_r - shoulder_x_l) ** 2 + (shoulder_y_r - shoulder_y_l) ** 2)
            #bili = delta_x_shoulder / delta_y_hiptoshoulder
            #bili_values.append(bili)
            #bili_mean = np.mean(bili_values)
            #bili_std = np.std(bili_values)
            # Write the closest angle to the output file
            closest_angle_label = "Right Arm Angle" if angle_in_degrees_r == closest_angle else "Left Arm Angle"
            f.write(
                f"Image: {filename},Closest to 90 Degrees: {closest_angle_label}, Angle: {closest_angle}")
            #,f" degrees, deltay: {bili}\n")

        print("Angles have been calculated and written to", output_file, closest_angle)
        #print(bili_mean, bili_std)
        mean = statistics.mean(closest_angles)
        print(mean)
